/*
 * Copyright Cedar Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! This module is as straightforward a translation as possible of
//! <https://github.com/cedar-policy/cedar-spec/blob/main/cedar-lean/Cedar/SymCC.lean>.

pub(crate) mod authorizer;
pub mod bitvec;
pub(crate) mod compiler;
pub(crate) mod concretizer;
mod decoder;
mod encoder;
pub(crate) mod enforcer;
mod entity_tag;
mod env;
pub mod ext;
pub mod extension_types;
mod extfun;
mod extractor;
pub mod factory;
mod function;
mod interpretation;
pub mod op;
mod result;
mod smtlib_script;
pub mod solver;
mod symbolizer;
mod tags;
pub mod term;
pub mod term_type;
pub mod type_abbrevs;
pub mod verifier;

use cedar_policy::Schema;
use cedar_policy_core::ast::{Expr, ExprBuilder, Policy, PolicySet};
use cedar_policy_core::validator::{
    typecheck::Typechecker, types::RequestEnv, ValidationMode, Validator,
};
use decoder::{parse_sexpr, IdMaps};
use encoder::Encoder;
use env::to_validator_request_env;
use solver::{Decision, Solver};
use std::borrow::Borrow;
use verifier::Asserts;

use crate::err::{Error, Result};

pub use bitvec::BitVecError;
pub use concretizer::ConcretizeError;
pub use concretizer::Env;
pub use decoder::DecodeError;
pub use encoder::EncodeError;
pub use env::{Environment, SymEnv};
pub use interpretation::Interpretation;
pub use result::CompileError;
pub use smtlib_script::SmtLibScript;
pub use solver::SolverError;
pub use verifier::{
    verify_always_allows, verify_always_denies, verify_always_matches, verify_disjoint,
    verify_equivalent, verify_implies, verify_never_errors, verify_never_matches,
};

/// Internal symbolic compiler.
/// Corresponds to the `SolverM` monad in Lean.
#[derive(Clone, Debug)]
pub struct SymCompiler<S> {
    solver: S,
}

impl<S> SymCompiler<S> {
    pub fn new(solver: S) -> Self {
        Self { solver }
    }

    pub fn solver(&self) -> &S {
        &self.solver
    }

    pub fn solver_mut(&mut self) -> &mut S {
        &mut self.solver
    }
}

impl<S: Solver> SymCompiler<S> {
    /// Given some `asserts` and their corresponding symbolic environment
    /// `symenv`, calls the SMT solver (if necessary) on an SMTLib encoding of
    /// `asserts` and returns `true` iff the result is unsatisfiable. The
    /// `asserts` are expected to be well-formed with respect to `symenv`. This
    /// call resets the solver.
    pub async fn check_unsat_asserts(
        &mut self,
        asserts: &Asserts,
        symenv: &SymEnv,
    ) -> Result<bool> {
        if asserts.iter().any(|assert| *assert == false.into()) {
            // some assert has been compiled to constant-false by symcc.
            // skip encoding and calling the solver.
            Ok(true)
        } else if asserts.iter().all(|assert| *assert == true.into()) {
            // all asserts have been compiled to constant-true by symcc.
            // skip encoding and calling the solver.
            Ok(false)
        } else {
            self.solver
                .smtlib_input()
                .reset()
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_logic("ALL")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            let mut encoder = Encoder::new(symenv, self.solver.smtlib_input())?;
            encoder.encode(asserts.iter()).await?;
            match self.solver.check_sat().await? {
                Decision::Unsat => Ok(true),
                Decision::Sat => Ok(false),
                Decision::Unknown => Err(Error::SolverUnknown),
            }
        }
    }

    /// Given a verification condition generator `vc` and a symbolic environment
    /// `symenv`, calls the SMT solver (if necessary) on an SMTLib encoding of
    /// the conditions generated by `vc` and returns `true` if the result is
    /// unsatisfiable. The function `vc` is expected to produce a list of terms
    /// type `Bool` that are well-formed with respect to the `symenv`. This call
    /// resets the solver.
    pub async fn check_unsat(
        &mut self,
        vc: impl FnOnce(&SymEnv) -> std::result::Result<Asserts, CompileError>,
        symenv: &SymEnv,
    ) -> Result<bool> {
        let asserts = vc(symenv)?;
        self.check_unsat_asserts(&asserts, symenv).await
    }

    /// Given some `asserts` and their corresponding symbolic environment `symenv`,
    /// calls the SMT solver (if necessary) on an SMTLib encoding of `asserts` and
    /// returns `None` if the result is unsatisfiable. Otherwise returns `Some` with
    /// a counterexample interpretation. The `asserts` are expected to be well-formed
    /// with respect to `symenv`. This call resets the solver.
    ///
    /// For soundness, `footprint` must include all expressions used for generating `asserts`.
    ///
    /// This function corresponds to `satAsserts?` (not `checkSatAsserts`) in the Lean.
    /// The `footprint` argument here corresponds to the `ps` argument in the Lean.
    pub async fn check_sat_asserts<E: Borrow<Expr>>(
        &mut self,
        asserts: &Asserts,
        symenv: &SymEnv,
        footprint: impl IntoIterator<Item = E>,
    ) -> Result<Option<Env>> {
        if asserts.iter().any(|assert| *assert == false.into()) {
            // some assert has been reduced to constant-false by the symcc process.
            // skip encoding and calling the solver.
            Ok(None)
        } else if asserts.iter().all(|assert| *assert == true.into()) {
            let interp = Interpretation::default(symenv);
            Ok(Some(symenv.interpret(&interp).concretize(footprint)?))
        } else {
            self.solver
                .smtlib_input()
                .reset()
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_option("produce-models", "true")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_logic("ALL")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            let mut encoder =
                Encoder::new(symenv, self.solver.smtlib_input()).map_err(Error::EncodeError)?;
            encoder
                .encode(asserts.iter())
                .await
                .map_err(Error::EncodeError)?;
            let id_maps = IdMaps::from_encoder(&encoder);
            match self.solver.check_sat().await? {
                Decision::Unsat => Ok(None),
                Decision::Sat => {
                    let Some(model_str) = self.solver.get_model().await? else {
                        return Ok(None);
                    };
                    let model = parse_sexpr(model_str.as_bytes())?;
                    let interp = model.decode_model(symenv, &id_maps)?;
                    Ok(Some(symenv.extract(footprint, &interp)?))
                }
                Decision::Unknown => Err(Error::SolverUnknown),
            }
        }
    }

    /// Returns true iff `policy` does not error on any well-formed input in the `symenv`.
    pub async fn check_never_errors(&mut self, policy: &Policy, symenv: &SymEnv) -> Result<bool> {
        self.check_unsat(|symenv| verify_never_errors(policy, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_never_errors`] is false.
    ///
    /// Corresponds to `neverErrors?` in the Lean.
    pub async fn check_never_errors_with_counterexample(
        &mut self,
        policy: &Policy,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        self.check_sat_asserts(
            &verify_never_errors(policy, symenv)?,
            symenv,
            std::iter::once(&policy.condition()),
        )
        .await
    }

    /// Returns true iff `policy` matches all well-formed inputs in the
    /// `symenv`.  That is, if `policy` is a `permit` policy, it allows all
    /// inputs in the `symenv`, or if `policy` is a `forbid` policy, it denies
    /// all inputs in the `symenv`.
    ///
    /// Compare with `check_always_allows`, which takes a policyset (which could
    /// consist of a single policy, or more) and determines whether it _allows_
    /// all well-formed inputs in a `symenv`. This function differs from
    /// `check_always_allows` on a singleton policyset in how it treats `forbid`
    /// policies -- while `check_always_allows` trivially doesn't hold for any
    /// policyset containing only `forbid` policies, `check_always_matches` does
    /// hold if the `forbid` policy explicitly denies all inputs in the
    /// `symenv`.
    pub async fn check_always_matches(&mut self, policy: &Policy, symenv: &SymEnv) -> Result<bool> {
        self.check_unsat(|symenv| verify_always_matches(policy, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_always_matches`] is false.
    ///
    /// Corresponds to `alwaysMatches?` in the Lean.
    pub async fn check_always_matches_with_counterexample(
        &mut self,
        policy: &Policy,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        self.check_sat_asserts(
            &verify_always_matches(policy, symenv)?,
            symenv,
            std::iter::once(&policy.condition()),
        )
        .await
    }

    /// Returns true iff `policy` matches no well-formed inputs in the
    /// `symenv`.
    ///
    /// Compare with `check_always_denies`, which takes a policyset (which could
    /// consist of a single policy, or more) and determines whether it _denies_
    /// all well-formed inputs in a `symenv`. This function differs from
    /// `check_always_denies` on a singleton policyset in how it treats `forbid`
    /// policies -- while `check_always_denies` trivially holds for any
    /// policyset containing only `forbid` policies, `check_never_matches` only
    /// holds if the `forbid` policy explicitly denies no inputs in the
    /// `symenv`.
    pub async fn check_never_matches(&mut self, policy: &Policy, symenv: &SymEnv) -> Result<bool> {
        self.check_unsat(|symenv| verify_never_matches(policy, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_never_matches`] is false.
    ///
    /// Corresponds to `neverMatches?` in the Lean.
    pub async fn check_never_matches_with_counterexample(
        &mut self,
        policy: &Policy,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        self.check_sat_asserts(
            &verify_never_matches(policy, symenv)?,
            symenv,
            std::iter::once(&policy.condition()),
        )
        .await
    }

    /// Returns true iff the authorization decision of `policies1` implies that
    /// of `policies2` for every well-formed input in the `symenv`. That is,
    /// every input allowed by `policies1` is allowed by `policies2`;
    /// `policies2` is either more permissive than, or equivalent to, `policies1`.
    pub async fn check_implies(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_implies(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_implies`] is false.
    ///
    /// Corresponds to `implies?` in the Lean.
    pub async fn check_implies_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition());
        self.check_sat_asserts(
            &verify_implies(policies1, policies2, symenv)?,
            symenv,
            footprint,
        )
        .await
    }

    /// Returns true iff `policies` allows all well-formed inputs in the `symenv`.
    pub async fn check_always_allows(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(|symenv| verify_always_allows(policies, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_always_allows`] is false.
    ///
    /// Corresponds to `alwaysAllows?` in the Lean.
    pub async fn check_always_allows_with_counterexample(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies
            .policies()
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat_asserts(
            &verify_always_allows(policies, symenv)?,
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff `policies` denies all well-formed inputs in the `symenv`.
    pub async fn check_always_denies(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(|symenv| verify_always_denies(policies, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_always_denies`] is false.
    ///
    /// Corresponds to `alwaysDenies?` in the Lean.
    pub async fn check_always_denies_with_counterexample(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies
            .policies()
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat_asserts(
            &verify_always_denies(policies, symenv)?,
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff `policies1` and `policies2` produce the same
    /// authorization decision on all well-formed inputs in the `symenv`.
    pub async fn check_equivalent(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_equivalent(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_equivalent`] is false.
    ///
    /// Corresponds to `equivalent?` in the Lean.
    pub async fn check_equivalent_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat_asserts(
            &verify_equivalent(policies1, policies2, symenv)?,
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff there is no well-formed input in the `symenv` that is allowed by both
    /// `policies1` and `policies2`. If this returns `false`, then there is at least one well-formed
    /// input that is allowed by both `policies1` and `policies2`.
    pub async fn check_disjoint(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_disjoint(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_disjoint`] is false.
    ///
    /// Corresponds to `disjoint?` in the Lean.
    pub async fn check_disjoint_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat_asserts(
            &verify_disjoint(policies1, policies2, symenv)?,
            symenv,
            footprint.iter(),
        )
        .await
    }
}

/// The Cedar symbolic compiler assumes that it receives well-typed policies.  This
/// function enforces this requirement by calling Cedar's typechecker. Specifically,
/// given a policy `p` and type environment `Γ`, this function calls the Cedar
/// typechecker to obtain a policy `p'` that is semantically equivalent to `p` and
/// well-typed with respect to `Γ`.
///
/// All `check` functions defined in this file _must_ be called on the _output_ of
/// this function (or `well_typed_policies`) to ensure that symbolic compilation
/// succeeds. Applying the symbolic compiler directly to a policy `p` may result in
/// type errors---that is, the compiler rejecting the policy because it does not
/// satisfy the `WellTyped` constraints that are assumed by the compiler, and
/// enforced by the typechecker through policy transformation.
pub fn well_typed_policy(
    policy: &Policy,
    env: &cedar_policy::RequestEnv,
    schema: &Schema,
) -> Result<Policy> {
    let env = to_validator_request_env(env, schema.as_ref())
        .ok_or_else(|| Error::ActionNotInSchema(env.action().to_string()))?;
    well_typed_policy_inner(policy, &env, schema)
}

fn well_typed_policy_inner(
    policy: &Policy,
    env: &RequestEnv<'_>,
    schema: &Schema,
) -> Result<Policy> {
    let validator_schema = schema.as_ref();
    // We need to perform these three checks like what the validator does here: https://github.com/cedar-policy/cedar/blob/82784864c01b5096cb73885dd2df5643074355ed/cedar-policy-core/src/validator.rs#L178-L185
    // We don't need `validate_template_action_application` because existence of `env` already serves as evidence
    let errs: Vec<_> =
        Validator::validate_entity_types_and_literals(schema.as_ref(), policy.template()).collect();
    if !errs.is_empty() {
        return Err(Error::PolicyNotWellTyped { errs });
    }
    let type_checker = Typechecker::new(validator_schema, ValidationMode::Strict);
    let policy_check = type_checker.typecheck_by_single_request_env(policy.template(), env);

    use cedar_policy_core::validator::typecheck::PolicyCheck::*;
    match policy_check {
        Success(expr) => Ok(Policy::from_when_clause(
            policy.effect(),
            expr.into_expr::<ExprBuilder<()>>(),
            policy.id().clone(),
            policy.loc().cloned(),
        )),
        Irrelevant(errs, expr) =>
        // A policy could be irrelevant just for this environment, so unless there were errors we don't want to fail.
        // Note that if the policy was irrelevant for all environments schema validation would have caught this
        // before SymCC. The Lean implementation needs to be updated to match this behavior.
        {
            if errs.is_empty() {
                Ok(Policy::from_when_clause(
                    policy.effect(),
                    expr.into_expr::<ExprBuilder<()>>(),
                    policy.id().clone(),
                    policy.loc().cloned(),
                ))
            } else {
                Err(Error::PolicyNotWellTyped { errs })
            }
        }
        Fail(errs) => Err(Error::PolicyNotWellTyped { errs }),
    }
}

/// The Cedar symbolic compiler assumes that it receives well-typed policies.  This
/// function enforces this requirement by calling Cedar's typechecker. Specifically,
/// given policies `ps` and a type environment `Γ`, this function calls the Cedar
/// typechecker on each `p ∈ ps` to obtain a policy `p'` that is semantically
/// equivalent to `p` and well-typed with respect to `Γ`.
///
/// All functions defined in this file that accept policies as input _must_ be
/// called on the _output_ of this function (or `wellTypedPolicy`) to ensure that
/// symbolic compilation succeeds. Applying the symbolic compiler directly to
/// policies `ps` may result in type errors---that is, the compiler rejecting the
/// policies because they don't satisfy the `WellTyped` constraints that are assumed
/// by the compiler, and enforced by the typechecker through policy transformation.
pub fn well_typed_policies(
    policies: &PolicySet,
    env: &cedar_policy::RequestEnv,
    schema: &Schema,
) -> Result<PolicySet> {
    let env = to_validator_request_env(env, schema.as_ref())
        .ok_or_else(|| Error::ActionNotInSchema(env.action().to_string()))?;
    let typed_policies: Result<Vec<Policy>> = policies
        .static_policies()
        .map(|p| well_typed_policy_inner(p, &env, schema))
        .collect();
    match typed_policies {
        Ok(ps) => {
            let mut res = PolicySet::new();
            // PANIC SAFETY
            #[allow(
                clippy::unwrap_used,
                reason = "adding well-typed policy should not error"
            )]
            ps.into_iter().for_each(|p| res.add(p).unwrap());
            Ok(res)
        }
        Err(err) => Err(err),
    }
}
